#include "compiler/parser.h"
#include <iostream>
#include <algorithm>
#include <string>

Parser::Parser(const std::vector<Token>& tokens)
    : tokens(tokens), current(0) {
}

std::unique_ptr<AST::Statement> Parser::parse() {
    return parseStatement();
}

std::vector<std::unique_ptr<AST::Statement>> Parser::parseMultipleStatements() {
    std::vector<std::unique_ptr<AST::Statement>> statements;
    
    while (!check(TokenType::END)) {
        // 尝试跳过任何前导的分号
        while (check(TokenType::DELIMITER, ";")) {
            advance();
        }
        
        if (check(TokenType::END)) break;
        
        // 解析一条语句
        statements.push_back(parseStatement());
        
        // 尝试跳过语句末尾的分号
        if (check(TokenType::DELIMITER, ";")) {
            advance();
        }
    }
    
    return statements;
}

std::unique_ptr<AST::Statement> Parser::parseStatement() {
    Token token = peek();

    if (check(TokenType::KEYWORD, "CREATE")) {
        advance();
        return parseCreateTable();
    }
    else if (check(TokenType::KEYWORD, "INSERT")) {
        advance();
        return parseInsert();
    }
    else if (check(TokenType::KEYWORD, "SELECT")) {
        advance();
        return parseSelect();
    }
    else if (check(TokenType::KEYWORD, "DELETE")) {
        advance();
        return parseDelete();
    }
    // Add handling for UPDATE in parseStatement method
    else if (check(TokenType::KEYWORD, "UPDATE")) {
        advance();
        return parseUpdate();
    }
    
    throw ParserError("Unexpected token: " + token.value, token.line, token.column);
}

// Move parseUpdate() method here as a member function implementation of Parser class
std::unique_ptr<AST::UpdateStmt> Parser::parseUpdate() {
    auto stmt = std::make_unique<AST::UpdateStmt>();

    Token tableName = expect(TokenType::IDENTIFIER, "Expected table name");
    stmt->tableName = tableName.value;

    expect(TokenType::KEYWORD, "SET", "Expected SET after table name");

    // 解析列名和新值对
    while (!check(TokenType::KEYWORD, "WHERE") && !check(TokenType::END)) {
        Token colName = expect(TokenType::IDENTIFIER, "Expected column name");
        expect(TokenType::OPERATOR, "=", "Expected '=' after column name");
        auto value = parseExpression();
        
        stmt->columnValues.emplace_back(colName.value, std::move(value));

        if (check(TokenType::DELIMITER, ",")) {
            advance();
        }
        else {
            break;
        }
    }

    if (check(TokenType::KEYWORD, "WHERE")) {
        advance();
        stmt->whereClause = parseExpression();
    }

    return stmt;
}

std::unique_ptr<AST::CreateTableStmt> Parser::parseCreateTable() {
    // 获取CREATE关键字的位置作为语句位置
    int line = tokens[current-1].line;
    int column = tokens[current-1].column;
    auto stmt = std::make_unique<AST::CreateTableStmt>(line, column);

    expect(TokenType::KEYWORD, "TABLE", "Expected TABLE after CREATE");

    Token tableName = expect(TokenType::IDENTIFIER, "Expected table name");
    stmt->tableName = tableName.value;

    expect(TokenType::DELIMITER, "(", "Expected '(' after table name");

    while (!check(TokenType::DELIMITER, ")") && !check(TokenType::END)) {
        Token colName = expect(TokenType::IDENTIFIER, "Expected column name");
        Token colType = expect(TokenType::KEYWORD, "Expected column type");

        stmt->columns.emplace_back(colName.value, colType.value);

        if (check(TokenType::DELIMITER, ",")) {
            advance();
        }
        else {
            break;
        }
    }

    expect(TokenType::DELIMITER, ")", "Expected ')' after column definitions");

    return stmt;
}

std::unique_ptr<AST::InsertStmt> Parser::parseInsert() {
    // 获取INSERT关键字的位置作为语句位置
    int line = tokens[current-1].line;
    int column = tokens[current-1].column;
    auto stmt = std::make_unique<AST::InsertStmt>(line, column);

    expect(TokenType::KEYWORD, "INTO", "Expected INTO after INSERT");

    Token tableName = expect(TokenType::IDENTIFIER, "Expected table name");
    stmt->tableName = tableName.value;

    if (check(TokenType::DELIMITER, "(")) {
        advance();

        while (!check(TokenType::DELIMITER, ")") && !check(TokenType::END)) {
            Token colName = expect(TokenType::IDENTIFIER, "Expected column name");
            stmt->columns.push_back(colName.value);

            if (check(TokenType::DELIMITER, ",")) {
                advance();
            }
            else {
                break;
            }
        }

        expect(TokenType::DELIMITER, ")", "Expected ')' after column list");
    }

    expect(TokenType::KEYWORD, "VALUES", "Expected VALUES after INSERT statement");
    expect(TokenType::DELIMITER, "(", "Expected '(' after VALUES");

    while (!check(TokenType::DELIMITER, ")") && !check(TokenType::END)) {
        auto value = parseExpression();
        stmt->values.push_back(std::move(value));

        if (check(TokenType::DELIMITER, ",")) {
            advance();
        }
        else {
            break;
        }
    }

    expect(TokenType::DELIMITER, ")", "Expected ')' after values");

    return stmt;
}

std::unique_ptr<AST::SelectStmt> Parser::parseSelect() {
    // 获取SELECT关键字的位置作为语句位置
    int line = tokens[current-1].line;
    int column = tokens[current-1].column;
    auto stmt = std::make_unique<AST::SelectStmt>(line, column);

    if (check(TokenType::OPERATOR, "*")) {
        advance();
        stmt->columns.push_back("*");
    }
    else {
        while (!check(TokenType::KEYWORD, "FROM") && !check(TokenType::END)) {
            Token colName = expect(TokenType::IDENTIFIER, "Expected column name");
            stmt->columns.push_back(colName.value);

            if (check(TokenType::DELIMITER, ",")) {
                advance();
            }
            else {
                break;
            }
        }
    }

    expect(TokenType::KEYWORD, "FROM", "Expected FROM after column list");

    Token tableName = expect(TokenType::IDENTIFIER, "Expected table name");
    stmt->tableName = tableName.value;

    if (check(TokenType::KEYWORD, "WHERE")) {
        advance();
        stmt->whereClause = parseExpression();
    }

    return stmt;
}

std::unique_ptr<AST::DeleteStmt> Parser::parseDelete() {
    // 获取DELETE关键字的位置作为语句位置
    int line = tokens[current-1].line;
    int column = tokens[current-1].column;
    auto stmt = std::make_unique<AST::DeleteStmt>(line, column);

    expect(TokenType::KEYWORD, "FROM", "Expected FROM after DELETE");

    Token tableName = expect(TokenType::IDENTIFIER, "Expected table name");
    stmt->tableName = tableName.value;

    if (check(TokenType::KEYWORD, "WHERE")) {
        advance();
        stmt->whereClause = parseExpression();
    }

    return stmt;
}

// Handle logical OR operator (lowest precedence)
std::unique_ptr<AST::Expression> Parser::parseExpression() {
    auto left = parseAndExpression();

    while (check(TokenType::KEYWORD, "OR")) {
        Token orToken = advance();
        auto right = parseAndExpression();
        left = std::make_unique<AST::BinaryExpression>(std::move(left), "OR", std::move(right), orToken.line, orToken.column);
    }

    return left;
}

// Handle logical AND operator (medium precedence)
std::unique_ptr<AST::Expression> Parser::parseAndExpression() {
    auto left = parseComparison();

    while (check(TokenType::KEYWORD, "AND")) {
        Token andToken = advance();
        auto right = parseComparison();
        left = std::make_unique<AST::BinaryExpression>(std::move(left), "AND", std::move(right), andToken.line, andToken.column);
    }

    return left;
}

// Handle comparison operators (higher precedence)
std::unique_ptr<AST::Expression> Parser::parseComparison() {
    auto left = parsePrimary();

    // 检查是否是BETWEEN操作符
    if (check(TokenType::KEYWORD, "BETWEEN")) {
        Token betweenToken = advance(); // 消费BETWEEN关键字
        auto lowerBound = parsePrimary(); // 解析下界值
        
        // 确保后面有AND关键字
        Token andToken = expect(TokenType::KEYWORD, "AND", "Expected AND after BETWEEN lower bound");
        
        auto upperBound = parsePrimary(); // 解析上界值
        
        // 获取标识符的位置信息
        int idLine = 0, idColumn = 0;
        if (dynamic_cast<AST::Identifier*>(left.get())) {
            idLine = left->line;
            idColumn = left->column;
        }
        
        // 对于BETWEEN操作，我们将其转换为两个比较操作的组合：column >= lower AND column <= upper
        auto geExpr = std::make_unique<AST::BinaryExpression>(
            std::make_unique<AST::Identifier>(dynamic_cast<AST::Identifier*>(left.get())->name, idLine, idColumn), 
            ">=", 
            std::move(lowerBound),
            betweenToken.line, betweenToken.column
        );
        
        auto leExpr = std::make_unique<AST::BinaryExpression>(
            std::make_unique<AST::Identifier>(dynamic_cast<AST::Identifier*>(left.get())->name, idLine, idColumn), 
            "<=", 
            std::move(upperBound),
            betweenToken.line, betweenToken.column
        );
        
        return std::make_unique<AST::BinaryExpression>(std::move(geExpr), "AND", std::move(leExpr), andToken.line, andToken.column);
    }

    if (check(TokenType::OPERATOR)) {
        Token op = advance();
        auto right = parsePrimary();
        return std::make_unique<AST::BinaryExpression>(std::move(left), op.value, std::move(right), op.line, op.column);
    }

    return left;
}

std::unique_ptr<AST::Expression> Parser::parsePrimary() {
    // Handle parenthetical expressions
    if (check(TokenType::DELIMITER, "(")) {
        Token leftParen = advance();
        auto expr = parseExpression();
        expect(TokenType::DELIMITER, ")", "Expected ')' after expression");
        // 保留内部表达式的位置信息
        return expr;
    }
    
    Token token = advance();

    if (token.type == TokenType::IDENTIFIER) {
        return std::make_unique<AST::Identifier>(token.value, token.line, token.column);
    }
    else if (token.type == TokenType::CONST_INT || token.type == TokenType::CONST_STRING) {
        return std::make_unique<AST::Literal>(token.value, token.type, token.line, token.column);
    }

    throw ParserError("Unexpected token in expression: " + token.value, token.line, token.column);
}

Token Parser::peek() const {
    if (current >= tokens.size()) return Token(TokenType::END, "", 0, 0);
    return tokens[current];
}

Token Parser::advance() {
    if (current >= tokens.size()) return Token(TokenType::END, "", 0, 0);
    return tokens[current++];
}

bool Parser::match(TokenType type) {
    if (check(type)) {
        advance();
        return true;
    }
    return false;
}

bool Parser::match(TokenType type, const std::string& value) {
    if (check(type, value)) {
        advance();
        return true;
    }
    return false;
}

Token Parser::expect(TokenType type, const std::string& message) {
    if (current >= tokens.size()) {
        throw ParserError(message + " but reached end of input", 0, 0);
    }

    Token token = tokens[current];
    if (token.type != type) {
        std::string expectedType;
        switch(type) {
            case TokenType::KEYWORD: expectedType = "KEYWORD";
                break;
            case TokenType::IDENTIFIER: expectedType = "IDENTIFIER";
                break;
            case TokenType::CONST_INT: expectedType = "NUMBER";
                break;
            case TokenType::CONST_STRING: expectedType = "STRING";
                break;
            case TokenType::OPERATOR: expectedType = "OPERATOR";
                break;
            case TokenType::DELIMITER: expectedType = "DELIMITER";
                break;
            default: expectedType = "UNKNOWN";
        }
        throw ParserError(message + ". Found '" + token.value + "' instead", token.line, token.column);
    }

    return advance();
}

// Helper function: Compare strings case-insensitively
bool Parser::equalsIgnoreCase(const std::string& a, const std::string& b) const {
    if (a.length() != b.length()) return false;
    
    for (size_t i = 0; i < a.length(); ++i) {
        if (std::toupper(a[i]) != std::toupper(b[i])) {
            return false;
        }
    }
    
    return true;
}

Token Parser::expect(TokenType type, const std::string& value, const std::string& message) {
    if (current >= tokens.size()) {
        throw ParserError(message + " but reached end of input", 0, 0);
    }

    Token token = tokens[current];
    if (token.type != type || !equalsIgnoreCase(token.value, value)) {
        throw ParserError(message + ". Expected '" + value + "', found '" + token.value + "' instead", token.line, token.column);
    }

    return advance();
}

bool Parser::check(TokenType type) const {
    if (current >= tokens.size()) return type == TokenType::END;
    return tokens[current].type == type;
}

bool Parser::check(TokenType type, const std::string& value) const {
    if (current >= tokens.size()) return false;
    return tokens[current].type == type && equalsIgnoreCase(tokens[current].value, value);
}